#!/bin/bash
#
# Clone from the NVMe disk into a image released bz2 archive
#

DISKDEV=
ROOTFSIMG="./system.img.raw"
ROOTFSIMG_SIZE=32
JetPackVer_Board=

MNTPOINT_DISK=/mnt/nvme_disk
MNTPOINT_IMG=/mnt/sys_image
JetPackVer=
L4TVer=
Board=
RKVER=
RELEASE_PKG=

OPTSTR="d:f:s:r:h"

function usage()
{
	echo -e "Usage:"
	echo -e "  $0 <opts>"
	echo -e "    -d <NVMeDisk> --------- The NVMe Disk device node in the host, can be: sdb, sdc, ..."
	echo -e "    -f <RootfsImg> -------- The rootfs image file path, default is: ${ROOTFSIMG}"
	echo -e "    -s <RootfsImg Size> --- The rootfs image file size, default is: ${ROOTFSIMG_SIZE}GB"
	echo -e "    -r <JetPackVer-Board> - Release the NVMe disk image(system.img.raw) into a bz2 archive"
	echo -e "                            The given JetPackVer_Board: 4.6.1-RKNS02NX0"
	echo -e "                            The package name format: Ritrontek_JetPack4.6.1_L4T32.7.1_RKNS02NX0_xxxxxx_rootfs_system.img.raw.tbz2"
	echo -e "    -h -------------------- Show this usage"
	echo -e "Examples:"
	echo -e "  $0 -d <NVMeDisk> [-f <RootfsImg>]"
	echo -e "  $0 -d <NVMeDisk> -r <JetPackVer-Board> [-f <RootfsImg>]"
}

function get_option()
{
	ARGS=$(getopt -o ${OPTSTR} -n "$0" -- "$@")
	if [[ $? != 0 ]];
	then
		usage
		exit 1
	fi

	eval set -- ${ARGS}

	while true; do
		case "$1" in
		-d)
			DISKDEV=$2
			shift 2
			;;
		-f)
			ROOTFSIMG=$2
			shift 2
			;;
		-s)
			ROOTFSIMG_SIZE=$2
			shift 2
			;;
		-r)
			JETPACKVER_BOARD=$2
			shift 2
			;;
		-h)
			usage
			exit 0
			;;
		--)
			break
			;;
		*)
			echo "[ERROR] Unsupport options: $2"
			usage
			exit 1
			;;
		esac
	done

	if [[ -z ${DISKDEV} ]]; then
		echo -e "[ERROR] NVMeDEV must be specified.\n"
		usage
		exit 1
	else
		if [[ ! ${DISKDEV} =~ ^sd[b-z]$ ]]; then
			echo -e "[ERROR] NVMe device name error, should be \"sdb, sdc, ...\".\n"
			exit 1
		fi
		if [[ ! -b /dev/${DISKDEV} ]]; then
			echo -e "[ERROR] NVMe device node \"/dev/${DISKDEV}\" does not exist.\n"
			exit 1
		fi
	fi

	if [[ ! ${JETPACKVER_BOARD} =~ [0-9]\.[0-9]\.[0-9]-RK[A-Z0-9]* ]]; then
		echo -e "[ERROR] JetPackVer_Board format error: ${JETPACKVER_BOARD}.\n"
		usage
		exit 1
	fi
	JetPackVer=$(echo ${JETPACKVER_BOARD} | awk -F "-" '{print $1}')
	if [[ ${JetPackVer} == "4.6.1" ]]; then
		L4TVer=32.7.1
	elif [[ ${JetPackVer} == "5.1.1" ]]; then
		L4TVer=35.3.1
	elif [[ ${JetPackVer} == "5.1.3" ]]; then
		L4TVer=35.5.0
	else
		echo -e "[ERROR] Unsupport JetPack Version: ${JetPackVer}.\n"
		usage
		exit 1
	fi

	Board=$(echo ${JETPACKVER_BOARD} | awk -F "-" '{print $2}')
}

function generate_blank_image()
{
	if [[ -f ${ROOTFSIMG} ]]; then
		echo -e "[INFO] Remove the existing image file"
		sudo rm ${ROOTFSIMG}
	fi
	echo -e "[INFO] Generate the blank image with \"dd\" ..."
	dd if=/dev/zero of=${ROOTFSIMG} bs=${ROOTFSIMG_SIZE}M count=1024
	echo -e "[INFO] Format to ext4 file system ..."
	sudo mkfs.ext4 ${ROOTFSIMG}
}

function copy_from_nvme_to_image()
{
	if [[ -d ${MNTPOINT_DISK} ]]; then
		local mountlist=$(findmnt -n ${MNTPOINT_DISK} | wc -l)
		if [[ ${mountlist} -gt 0 ]]; then
			sudo umount ${MNTPOINT_DISK}
		fi
	else
		sudo mkdir ${MNTPOINT_DISK}
	fi

	if [[ -d ${MNTPOINT_IMG} ]]; then
		mountlist=$(findmnt -n ${MNTPOINT_IMG} | wc -l)
		if [[ ${mountlist} -gt 0 ]]; then
			sudo umount ${MNTPOINT_IMG}
		fi
	else
		sudo mkdir ${MNTPOINT_IMG}
	fi

	# Copy rootfs files from nvme disk to image
	sudo mount /dev/${DISKDEV}1 ${MNTPOINT_DISK}
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Mount /dev/${DISKDEV}1 error.\n"
		exit 1
	fi
	sudo mount ${ROOTFSIMG} ${MNTPOINT_IMG}
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Mount ${ROOTFSIMG} error.\n"
		exit 1
	fi

	sudo rsync -axHAWX --numeric-ids --info=progress2 --exclude=/proc ${MNTPOINT_DISK}/ ${MNTPOINT_IMG}

	sudo umount ${MNTPOINT_DISK} ${MNTPOINT_IMG}

	sudo e2fsck -p -f ${ROOTFSIMG}
	sudo resize2fs -M ${ROOTFSIMG}
	echo -e "[INFO] The finally ${ROOTFSIMG} infomation:"
	echo -e ">>> $(ls -hl ${ROOTFSIMG}).\n"
}

function generate_release_pkg_name()
{
	if [[ -d ${MNTPOINT_DISK} ]]; then
		local mountlist=$(findmnt -n ${MNTPOINT_DISK} | wc -l)
		if [[ ${mountlist} -gt 0 ]]; then
			sudo umount ${MNTPOINT_DISK}
		fi
	else
		sudo mkdir ${MNTPOINT_DISK}
	fi

	sudo mount /dev/${DISKDEV}1 ${MNTPOINT_DISK}
	RET=$?
	if [[ ${RET} -ne 0 ]]; then
		echo -e "[ERROR] Mount /dev/${DISKDEV}1 error.\n"
		exit 1
	fi

	RKVER=$(strings ${MNTPOINT_DISK}/boot/Image  | grep -i "Linux version [0-9]" | awk '{print $3}' | awk -F '-' '{print $3}')

	sudo umount ${MNTPOINT_DISK}

	if [[ -z ${JetPackVer} ]] \
		|| [[ -z ${L4TVer} ]] \
		|| [[ -z ${Board} ]] \
		|| [[ -z ${RKVER} ]];
	then
		echo -e "JetPackVer=${JetPackVer}"
		echo -e "L4TVer=${L4TVer}"
		echo -e "Board=${Board}"
		echo -e "RKVER=${RKVER}"
		echo -e "[ERROR] RELEASE_PKG name error.\n"
		exit 1
	fi

	RELEASE_PKG=Ritrontek_JetPack${JetPackVer}_L4T${L4TVer}_${Board}_${RKVER}_rootfs_system.img.raw.tbz2
}

function release_image_pkg()
{
	tar -jcvf ${RELEASE_PKG} ${ROOTFSIMG}
}

####################################################################

if [[ $(id -u) -ne 0 ]]; then
	echo -e "[ERROR] Please run as root!\n"
	exit 1
fi

get_option $@

echo "DISKDEV=${DISKDEV}"
echo "ROOTFSIMG=${ROOTFSIMG}"
echo "ROOTFSIMG_SIZE=${ROOTFSIMG_SIZE}"

echo -e "\n[INFO] Generate blank image file"
echo -e ">>> ${ROOTFSIMG} (${ROOTFSIMG_SIZE}GB)"
generate_blank_image

echo -e "\n[INFO] Copy rootfs files from DISKDEV to ROOTFSIMG ..."
echo -e ">>> From: /dev/${DISKDEV}1"
echo -e ">>> To  : ${ROOTFSIMG}"
copy_from_nvme_to_image

generate_release_pkg_name
if [[ -n ${RELEASE_PKG} ]]; then
	echo -e "\n[INFO] Release rootfs image into a archive"
	echo -e ">>> ${RELEASE_PKG}\n"
	release_image_pkg
fi

echo -e "\n[INFO] Done!"
